import math
import numpy as np
from spatialmath import SO3, SE3


class Calibration:
    @staticmethod
    def calibrate(Xbe: np.ndarray, Xco: np.ndarray, Xcw: np.ndarray):
        tcb, rpy_cb, tbc, rpy_bc, teo = Calibration.base_calibrate(Xbe, Xco)
        tcw, rpy_cw = Calibration.workpiece_calibrate(Xcw)

        Tcb: SE3 = SE3.Trans(*tcb) * SE3.RPY(*rpy_cb)
        Tcw: SE3 = SE3.Trans(*tcw) * SE3.RPY(*rpy_cw)

        Tbw: SE3 = Tcb.inv() * Tcw
        tbw, rpy_bw = Tbw.t, Tbw.rpy()
        return tbw, rpy_bw

    @staticmethod
    def base_calibrate(Xbe: np.ndarray, Xco: np.ndarray) -> tuple:
        # 前四个点求基座相对于视觉的姿态
        Xbe1 = Xbe[:4, :]
        Xco1 = Xco[:4, :]

        num1 = math.comb(4, 2)
        A1 = np.zeros((3, num1))
        b1 = np.zeros((3, num1))

        k = 0
        for i in range(4 - 1):
            for j in range(i + 1, 4):
                tco1 = Xco1[i, :]
                tco2 = Xco1[j, :]
                tbe1 = Xbe1[i, :3]
                tbe2 = Xbe1[j, :3]
                A1[:, k] = tco1 - tco2
                b1[:, k] = tbe1 - tbe2
                k += 1
        Rbc = b1 @ np.linalg.pinv(A1)
        U, S, Vh = np.linalg.svd(Rbc)
        Rbc = U @ Vh
        Rcb = np.transpose(Rbc)
        rpy_cb = SO3(Rcb).rpy().squeeze()
        rpy_bc = SO3(Rbc).rpy().squeeze()

        # 后四个点求基座相对于视觉的位置
        Xbe2 = Xbe[4:, :]
        Xco2 = Xco[4:, :]
        num2 = math.comb(Xbe2.shape[0], 2)
        A2 = np.zeros((3 * num2, 3))
        b2 = np.zeros((3 * num2, 1))

        k = 0
        for i in range(Xbe2.shape[0] - 1):
            for j in range(i + 1, Xbe2.shape[0]):
                Rbe1 = SO3.RPY(*Xbe2[i, 3:]).R
                Rbe2 = SO3.RPY(*Xbe2[j, 3:]).R
                tco1 = np.reshape(Xco2[i, :], (3, 1))
                tco2 = np.reshape(Xco2[j, :], (3, 1))
                tbe1 = np.reshape(Xbe2[i, :3], (3, 1))
                tbe2 = np.reshape(Xbe2[j, :3], (3, 1))
                A2[k * 3:(k + 1) * 3, :] = np.transpose(Rbe1 - Rbe2)
                b2[k * 3:(k + 1) * 3, :] = - (Rbe1.T @ Rbc @ tco1 - Rbe2.T @ Rbc @ tco2) \
                                           + (Rbe1.T @ tbe1 - Rbe2.T @ tbe2)
                k += 1

        tbc = np.linalg.pinv(A2) @ b2
        tcb = - Rcb @ tbc
        tbc = tbc.squeeze()
        tcb = tcb.squeeze()

        # 求工具相对于末端的位置
        teo = np.zeros(3)
        for i in range(Xbe.shape[0]):
            Rbe = SO3.RPY(*Xbe[i, 3:]).R
            tco = Xco[i, :]
            tbe = Xbe[i, :3]
            teo += Rbe.T @ Rbc @ tco + Rbe.T @ tbc - Rbe.T @ tbe
        teo /= Xbe.shape[0]

        return tcb, rpy_cb, tbc, rpy_bc, teo

    @staticmethod
    def workpiece_calibrate(Xcw: np.ndarray) -> tuple:
        op = np.squeeze(Xcw[0, :])
        ox = np.squeeze(Xcw[1, :] - Xcw[0, :])
        oy = np.squeeze(Xcw[2, :] - Xcw[1, :])

        ox = ox / np.linalg.norm(ox)
        oy = oy / np.linalg.norm(oy)

        oz = np.cross(ox, oy)
        oy = np.cross(oz, ox)

        T = np.eye(4)
        T[:3, :] = np.transpose(np.vstack((ox, oy, oz, op)))
        T = SE3(T)
        tcw, rpy_cw = T.t, T.rpy()
        return tcw, rpy_cw
